import numpy as np
from collections import deque
import matplotlib.pyplot as plt

class Logger:
    """训练日志记录器"""
    def __init__(self, log_interval=100):
        self.log_interval = log_interval
        self.episode_rewards = []
        self.episode_lengths = []
        self.losses = []
        self.q_values = []
        self.epsilons = []
        self.episode_count = 0
        self.recent_rewards = deque(maxlen=100)

    def log_episode(self, episode, reward, length, epsilon):
        self.episode_rewards.append(reward)
        self.episode_lengths.append(length)
        self.epsilons.append(epsilon)
        self.recent_rewards.append(reward)
        self.episode_count = episode

        if episode % self.log_interval == 0:
            avg_reward = np.mean(self.recent_rewards)
            print(f"Episode {episode}, "
                  f"Average Reward: {avg_reward:.2f}, "
                  f"Last Reward: {reward:.2f}, "
                  f"Epsilon: {epsilon:.3f}")

    def log_loss(self, loss, q_value):
        self.losses.append(loss)
        self.q_values.append(q_value)